import os
from typing import Dict, Any, Optional
from src.core.llm_config import ModelRouter
from src.agents.research.research_agents import (
    FundamentalResearchAgent as FundamentalAgent,
    IndustryResearchAgent as IndustryAgent,
    MacroResearchAgent as MacroAgent,
    MarketObserverAgent,
    RiskAgent,
    AdvisorAgent
)
from src.agents.base import AnalystAgent  # 保留原有的分析师Agent

class OrchestratorAgent:
    """主控Agent，负责协调各个专业Agent的工作"""
    
    def __init__(self):
        # 初始化模型路由
        self.model_router = ModelRouter()
        
        # 初始化各个专业Agent
        self.fundamental_agent = FundamentalAgent(self.model_router)
        self.industry_agent = IndustryAgent(self.model_router)
        self.macro_agent = MacroAgent(self.model_router)
        self.market_observer = MarketObserverAgent(self.model_router)
        self.risk_agent = RiskAgent(self.model_router)
        self.advisor_agent = AdvisorAgent(self.model_router)
        self.analyst_agent = AnalystAgent(self.model_router)  # 兼容原有功能
    
    def process_request(self, user_input: str, selected_role: str = "全流程") -> Dict[str, Any]:
        """
        处理用户请求，根据选择的角色调用相应的Agent
        
        Args:
            user_input: 用户输入的查询内容
            selected_role: 选择的角色（全流程/基本面研究/行业研究/宏观研究/市场观察/风控/顾问）
            
        Returns:
            包含分析结果的字典
        """
        results = {}
        
        try:
            if selected_role == "全流程":
                results = self._process_full_workflow(user_input)
            elif selected_role == "基本面研究":
                results = self._process_fundamental_research(user_input)
            elif selected_role == "行业研究":
                results = self._process_industry_research(user_input)
            elif selected_role == "宏观研究":
                results = self._process_macro_research()
            elif selected_role == "市场观察":
                results = self._process_market_observation(user_input)
            elif selected_role == "风控":
                results = self._process_risk_assessment(user_input)
            elif selected_role == "顾问":
                results = self._process_advice_generation(user_input)
            else:
                # 默认使用原有分析师Agent
                results = self._process_analyst_workflow(user_input)
                
        except Exception as e:
            results["error"] = f"处理请求时发生错误: {str(e)}"
            
        return results
    
    def _process_full_workflow(self, symbol: str) -> Dict[str, Any]:
        """处理全流程分析"""
        results = {}
        
        # 1. 基本面研究
        fundamental_result = self.fundamental_agent.analyze_fundamentals(symbol)
        results["fundamental"] = fundamental_result
        
        # 2. 行业研究
        # 从基本面分析中提取行业信息（简化实现）
        industry = f"{symbol}所属行业"
        industry_result = self.industry_agent.analyze_industry(industry)
        results["industry"] = industry_result
        
        # 3. 宏观研究
        macro_result = self.macro_agent.analyze_macro()
        results["macro"] = macro_result
        
        # 4. 市场观察
        market_result = self.market_observer.observe_market(symbol)
        results["market"] = market_result
        
        # 5. 风险评估
        risk_result = self.risk_agent.assess_risk(results)
        results["risk"] = risk_result
        
        # 6. 投资建议
        advice_result = self.advisor_agent.provide_advice(results)
        results["advice"] = advice_result
        
        return results
    
    def _process_fundamental_research(self, symbol: str) -> Dict[str, Any]:
        """处理基本面研究"""
        return {
            "fundamental": self.fundamental_agent.analyze_fundamentals(symbol)
        }
    
    def _process_industry_research(self, industry: str) -> Dict[str, Any]:
        """处理行业研究"""
        return {
            "industry": self.industry_agent.analyze_industry(industry)
        }
    
    def _process_macro_research(self) -> Dict[str, Any]:
        """处理宏观研究"""
        return {
            "macro": self.macro_agent.analyze_macro()
        }
    
    def _process_market_observation(self, symbol: Optional[str] = None) -> Dict[str, Any]:
        """处理市场观察"""
        return {
            "market": self.market_observer.observe_market(symbol)
        }
    
    def _process_risk_assessment(self, analysis_data: str) -> Dict[str, Any]:
        """处理风险评估"""
        # 简化实现，实际应用中需要更复杂的分析数据结构
        analysis_results = {"data": analysis_data}
        return {
            "risk": self.risk_agent.assess_risk(analysis_results)
        }
    
    def _process_advice_generation(self, analysis_data: str) -> Dict[str, Any]:
        """处理投资建议生成"""
        # 简化实现，实际应用中需要更复杂的分析数据结构
        analyses = {"fundamental": {"analysis": analysis_data}}
        return {
            "advice": self.advisor_agent.provide_advice(analyses)
        }
    
    def _process_analyst_workflow(self, query: str) -> Dict[str, Any]:
        """处理原有分析师工作流（兼容性）"""
        analysis_result = self.analyst_agent.analyze_market_trends(query)
        risk_result = self.risk_agent.assess_risk({"query": query, "analysis": analysis_result})
        advice_result = self.advisor_agent.provide_advice({
            "fundamental": {"analysis": analysis_result},
            "risk": risk_result
        })
        
        return {
            "analysis": analysis_result,
            "risk": risk_result,
            "advice": advice_result
        }

# MCP服务封装
def create_mcp_server():
    """
    创建MCP服务服务器，将Agent功能暴露为工具
    """
    try:
        from fastmcp import FastMCP
        
        # 初始化FastMCP服务器
        mcp = FastMCP("financial_agent_server")
        
        # 初始化主控Agent
        orchestrator = OrchestratorAgent()
        
        # 注册工具
        @mcp.tool()
        def fundamental_analysis(symbol: str) -> dict:
            """公司基本面分析工具"""
            try:
                result = orchestrator.process_request(symbol, "基本面研究")
                return {"status": "success", "data": result}
            except Exception as e:
                return {"status": "error", "message": str(e)}
        
        @mcp.tool()
        def industry_analysis(industry: str) -> dict:
            """行业研究分析工具"""
            try:
                result = orchestrator.process_request(industry, "行业研究")
                return {"status": "success", "data": result}
            except Exception as e:
                return {"status": "error", "message": str(e)}
        
        @mcp.tool()
        def macro_analysis() -> dict:
            """宏观经济分析工具"""
            try:
                result = orchestrator.process_request("", "宏观研究")
                return {"status": "success", "data": result}
            except Exception as e:
                return {"status": "error", "message": str(e)}
        
        @mcp.tool()
        def market_observation(symbol: Optional[str] = None) -> dict:
            """市场观察工具"""
            try:
                result = orchestrator.process_request(symbol or "", "市场观察")
                return {"status": "success", "data": result}
            except Exception as e:
                return {"status": "error", "message": str(e)}
        
        @mcp.tool()
        def risk_assessment(analysis_data: str) -> dict:
            """投资风险评估工具"""
            try:
                result = orchestrator.process_request(analysis_data, "风控")
                return {"status": "success", "data": result}
            except Exception as e:
                return {"status": "error", "message": str(e)}
        
        @mcp.tool()
        def investment_advice(analysis_data: str) -> dict:
            """投资建议生成工具"""
            try:
                result = orchestrator.process_request(analysis_data, "顾问")
                return {"status": "success", "data": result}
            except Exception as e:
                return {"status": "error", "message": str(e)}
        
        @mcp.tool()
        def full_analysis(symbol: str) -> dict:
            """全流程金融分析工具"""
            try:
                result = orchestrator.process_request(symbol, "全流程")
                return {"status": "success", "data": result}
            except Exception as e:
                return {"status": "error", "message": str(e)}
        
        @mcp.tool()
        def health_check() -> dict:
            """健康检查端点"""
            return {
                "status": "healthy",
                "service": "financial-agent-orchestrator",
                "version": "1.0.0"
            }
        
        return mcp
        
    except ImportError:
        print("Warning: FastMCP not available, MCP server functionality disabled")
        return None
    except Exception as e:
        print(f"Error creating MCP server: {e}")
        return None

def start_mcp_server(host: str = "0.0.0.0", port: int = 8081):
    """
    启动MCP服务器
    
    Args:
        host: 服务器主机地址
        port: 服务器端口
    """
    mcp_server = create_mcp_server()
    if mcp_server:
        print(f"Starting Orchestrator MCP Server on {host}:{port}")
        try:
            mcp_server.run(host=host, port=port)
        except Exception as e:
            print(f"Failed to start MCP server: {e}")
    else:
        print("MCP server not available")

if __name__ == "__main__":
    # 创建并启动MCP服务器
    mcp_port = int(os.getenv("ORCHESTRATOR_MCP_PORT", 8081))
    mcp_host = os.getenv("ORCHESTRATOR_MCP_HOST", "0.0.0.0")
    
    print("Orchestrator Agent initialized successfully")
    start_mcp_server(host=mcp_host, port=mcp_port)